import yaml


def read_yaml_number_param(file_path, param_key):
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            data = yaml.safe_load(file)
            if param_key in data and isinstance(data[param_key], (int, float)):
                return data[param_key]
            else:
                print(f"未找到参数 {param_key} 或者该参数不是数字类型。")
                return None
    except FileNotFoundError:
        print(f"错误：文件 {file_path} 未找到。")
        return None
    except yaml.YAMLError as e:
        print(f"错误：解析 YAML 文件时出错 - {e}")
        return None


def write_yaml_number_param(file_path, param_key, new_value):
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            data = yaml.safe_load(file)

        if param_key in data and isinstance(data[param_key], (int, float)):
            data[param_key] = new_value
            with open(file_path, 'w', encoding='utf-8') as file:
                yaml.dump(data, file, default_flow_style=False)
            # print(f"参数 {param_key} 已更新为 {new_value}。")
        else:
            print(f"未找到参数 {param_key} 或者该参数不是数字类型。")
    except FileNotFoundError:
        print(f"错误：文件 {file_path} 未找到。")
    except yaml.YAMLError as e:
        print(f"错误：解析或写入 YAML 文件时出错 - {e}")


if __name__ == "__main__":
    file_path = 'bak/fail_1/epoch.yaml'
    param_key = 'epoch_num'

    # 读取参数
    current_value = read_yaml_number_param(file_path, param_key)
    if current_value is not None:
        print(f"当前参数 {param_key} 的值为: {current_value}")

        # 修改参数
        new_value = current_value + 1
        write_yaml_number_param(file_path, param_key, new_value)